import argparse

from Models.utils import activate_bn

def eval_config():
    '''
        common paramaters for semi and supervised test setting
    '''
    parser = argparse.ArgumentParser()
    # about task
    parser.add_argument('-way', type=int, default=5)
    parser.add_argument('-shot', type=int, default=1)
    # number of query image per class
    parser.add_argument('-query', type=int, default=15)
    parser.add_argument('-dataset', type=str, default='miniimagenet', choices=['miniimagenet', 'cub', 'tieredimagenet', 'fc100',
                                                                            'tieredimagenet_yao', 'cifar_fs', 'recognition36', 'recognition36_crop', 'cars', 'VHR-10'])  # 增加自定义数据集
    parser.add_argument('-set', type=str, default='test',
                        choices=['train', 'val', 'test'])
    # about model
    parser.add_argument('-temperature', type=float, default=12.5)
    parser.add_argument('-metric', type=str, default='cosine', choices=['cosine', 'l2'])
    parser.add_argument('-norm', type=str, default='center', choices=['center'])
    parser.add_argument('-deepemd', type=str, default='fcn',
                        choices=['fcn', 'grid', 'sampling'])
    # deepemd fcn only
    parser.add_argument('-feature_pyramid', type=str, default=None)
    # deepemd sampling only
    parser.add_argument('-num_patch', type=int, default=9)
    # deepemd grid only patch_list
    parser.add_argument('-patch_list', type=str, default='2,3')
    parser.add_argument('-patch_ratio', type=float, default=2)
    # solver
    parser.add_argument('-solver', type=str, default='opencv', choices=['opencv'])
    # SFC
    parser.add_argument('-sfc_lr', type=float, default=100.0)
    parser.add_argument('-sfc_wd', type=float, default=0,
                        help='weight decay for SFC weight')
    parser.add_argument('-sfc_update_step', type=int, default=100)
    parser.add_argument('-sfc_bs', type=int, default=4)
    # others
    parser.add_argument('-test_episode', type=int, default=1000)
    parser.add_argument('-gpu', default='0,1')
    parser.add_argument('-seed', type=int, default=1)
    parser.add_argument('-extra_dir', type=str, default=None,
                        help='extra information that is added to checkpoint dir, e.g. hyperparameters')
    parser.add_argument('--image_size', type=int, default=84,
                        help='extra information that is added to checkpoint dir, e.g. hyperparameters')
    # ====================================自定义模型参数====================================
    # 额外参数
    parser.add_argument("--origin", action="store_true", help="使用论文参数")
    parser.add_argument('--model', type=str, default='resnet',
                        help='选择要使用的backbone(为vit transformer做准备), 使用ViT作为backbone时请使用FCN模式')
    parser.add_argument("--model_name", type=str, default="resnet", help="用于选择使用的模型")
    parser.add_argument('--not_use_clstoken', action="store_true",
                        help='viT模型可选项是否添加cls token, 默认使用')
    parser.add_argument('--vit_mode', type=str, default='cls',
                        choices=['cls', 'mean'], help='选择使用cls token或者mean(平均所有patch)的方式')
    parser.add_argument('--vit_depth', type=int, default=4, help="使用ViT时的深度")
    parser.add_argument('--not_imagenet_pretrain',
                        action="store_true", help="是否使用imagenet的pretrain参数")
    loss = parser.add_mutually_exclusive_group()
    loss.add_argument('--use_taylor', action='store_true', help="使用通常版本的taylor softmax loss")
    loss.add_argument('--use_elu_taylor', action='store_true', help='使用elu taylor softmax loss')
    loss.add_argument('--use_sig_softmax', action='store_true', help="使用 sigsoftmax 作为损失函数")
    loss.add_argument('--use_focal_loss', action='store_true', help='使用 focal loss 作为损失函数')

    parser.add_argument('-uss', "--use_specific_status", action="store_true", help='是否使用数据集相关统计信息进行预处理中的normalize')

    # resnet下使用注意力机制的相关参数
    parser.add_argument('--with_SA', action='store_true', help="在resnet基础上使用self-attention模式")
    parser.add_argument('--multi_layer_atten', default=None, dest="MLA",help="在resnet多个层之间使用attention")
    parser.add_argument('--SA_heads', type=int, default=8, help="resnet使用heads的数目")
    parser.add_argument('--SA_mlp_dim', type=int, default=1024, help="resnet中SA模块使用的mlp中隐藏层的数目")
    parser.add_argument('--SA_depth', type=int, default=1, help='resnet下SA模块的层数')
    parser.add_argument('--SA_dim_head', type=int, default=64, help="resnet下SA模块每个head的维度")
    parser.add_argument('--SA_dropout', type=float, default=0.1, help="resnet下SA模块的dropout率")
    parser.add_argument('--SA_res', action="store_true", help="使用残差连接")
    parser.add_argument('--SA_relu', action="store_true", help="add relu after attention")
    parser.add_argument('--no_mlp', action="store_true", help="去除mlp层")
    parser.add_argument('--sche', type=str,default="StepLR" ,choices=['StepLR', 'CosineLR'], help="使用StepLR或者CosineAnnealingLR scheduler")
    attn_type_arg = parser.add_mutually_exclusive_group()
    attn_type_arg.add_argument('--pos_embed', action="store_true", help="加入相对位置编码(relative position embedding)")
    attn_type_arg.add_argument('--use_rpe', action="store_true", help="加入相对位置编码(使用rpe模块)")
    attn_type_arg.add_argument('--use_origin_outlook', action="store_true", help="使用Outlook版本的attn")
    attn_type_arg.add_argument('--use_involution', action="store_true", help="使用 involution 的注意力方法")
    attn_type_arg.add_argument('--use_carafe_attn', action="store_true", help='使用CARAFE的注意力方法')

    parser.add_argument("--outlook_kernel", type=int, default=3, help="outlook attention 使用的 Kernel size")
    parser.add_argument("--involution_softmax", action="store_true", help="involution 模块中是否使用Softmax机制")
    parser.add_argument("--involution_kernel", type=int, default=3, help="involution attention 使用的 Kernel size")
    parser.add_argument("--carafe_kernel", type=int, default=3, help="指定 carafe 模块 kernal_size ")

    parser.add_argument("--rpe_method", default="product" , type=str, choices=["euc", "quant", "cross", "product"])
    parser.add_argument("--rpe_ratio", default=1.9, type=float, help="制定RPE模块的Ratio")
    parser.add_argument("--rpe_mode", default="Contextual", type=str, choices=[None, "bias", "contextual"])
    parser.add_argument("--rpe_on",default="k", type=str, choices=["k", "qk", "qkv", "q", "v"])

    # resnet可变形卷积相关参数
    parser.add_argument("--use_deformconv", action='store_true', help="在resnet12基础上使用可变形卷积")
    parser.add_argument("--modulated_deformconv", action="store_true", help="在resnet12基础上使用可变形卷积v2")
    # ====================================自定义模型参数====================================
    return parser

def train_config(parser : argparse.ArgumentParser):
    '''
        common model parameters for pre-train and meta-train
    '''
    # ================================== 自定义的相关参数 ======================================
    parser.add_argument('--model', type=str, default='resnet',
                        help='选择要使用的backbone(为vit transformer做准备), 使用ViT作为backbone时请使用FCN模式')
    parser.add_argument('--not_use_clstoken', action="store_true",
                        help='viT模型可选项是否添加cls token, 默认使用')
    parser.add_argument('--vit_mode', type=str, default='cls',
                        choices=['cls', 'mean'], help='选择使用cls token或者mean(平均所有patch)的方式')
    parser.add_argument('--vit_depth', type=int, default=4, help="使用ViT时的深度")
    parser.add_argument('--not_imagenet_pretrain',
                        action="store_true", help="是否使用imagenet的pretrain参数")
    parser.add_argument('--sche', type=str,default="StepLR" ,choices=['StepLR', 'CosineLR'], help="使用StepLR或者CosineAnnealingLR scheduler")
    loss = parser.add_mutually_exclusive_group()
    loss.add_argument('--use_taylor', action='store_true', help="使用通常版本的taylor softmax loss")
    loss.add_argument('--use_elu_taylor', action='store_true', help='使用elu taylor softmax loss')
    loss.add_argument('--use_sig_softmax', action='store_true', help="使用 sigsoftmax 作为损失函数")
    loss.add_argument('--use_focal_loss', action='store_true', help='使用 focal loss 作为损失函数')

    # ========================= resnet下使用注意力机制的相关参数 =========================
    parser.add_argument('--with_SA', action='store_true', help="在resnet基础上使用self-attention模式")
    parser.add_argument('--multi_layer_atten', default=None, dest="MLA",help="在resnet多个层之间使用attention")
    parser.add_argument('--SA_heads', type=int, default=8, help="resnet使用heads的数目")
    parser.add_argument('--SA_mlp_dim', type=int, default=1024, help="resnet中SA模块使用的mlp中隐藏层的数目")
    parser.add_argument('--SA_depth', type=int, default=1, help='resnet下SA模块的层数')
    parser.add_argument('--SA_dim_head', type=int, default=64, help="resnet下SA模块每个head的维度")
    parser.add_argument('--SA_dropout', type=float, default=0.1, help="resnet下SA模块的dropout率")
    parser.add_argument('--SA_res', action="store_true", help="使用残差连接")
    parser.add_argument('--SA_relu', action="store_true", help="add relu after attention")
    parser.add_argument('--no_mlp', action="store_true", help="去除mlp层")
    
    attn_type_arg = parser.add_mutually_exclusive_group()
    attn_type_arg.add_argument('--pos_embed', action="store_true", help="加入相对位置编码(relative position embedding)")
    attn_type_arg.add_argument('--use_rpe', action="store_true", help="加入相对位置编码(使用rpe模块)")
    attn_type_arg.add_argument('--use_origin_outlook', action="store_true", help="使用Outlook版本的attn")
    attn_type_arg.add_argument('--use_involution', action="store_true", help="使用 involution 的注意力方法")
    attn_type_arg.add_argument('--use_carafe_attn', action="store_true", help='使用CARAFE的注意力方法')
    parser.add_argument("--outlook_kernel", type=int, default=3, help="outlook attention 使用的 Kernel size")
    parser.add_argument("--involution_softmax", action="store_true", help="involution 模块中是否使用Softmax机制")
    parser.add_argument("--involution_kernel", type=int, default=3, help="involution attention 使用的 Kernel size")
    parser.add_argument("--carafe_kernel", type=int, default=3, help="指定 carafe 模块 kernal_size ")

    parser.add_argument("--rpe_method", default="product" , type=str, choices=["euc", "quant", "cross", "product"])
    parser.add_argument("--rpe_ratio", default=1.9, type=float, help="制定RPE模块的Ratio")
    parser.add_argument("--rpe_mode", default="contextual", type=str, choices=[None, "bias", "contextual"])
    parser.add_argument("--rpe_on",default="k", type=str, choices=["k", "qk", "qkv", "q", "v"])

    # ================================== 可变形卷积相关参数 ===================================================
    parser.add_argument("--use_deformconv", action='store_true', help="在resnet12基础上使用可变形卷积")
    parser.add_argument("--modulated_deformconv", action="store_true", help="在resnet12基础上使用可变形卷积v2")
    return parser
